import astro.fit
import numpy as np
import matplotlib.pyplot as pl

def gauss2d(xy,x0,y0,xfwhm,yfwhm,height):
    """ Calculate a 2d gaussian."""
    x,y = xy
    const2 = 2*1.177410023               # 2*sqrt(2*ln(2))
    xsig = xfwhm / const2
    ysig = yfwhm / const2
    term1 = (x-x0)**2 / (2*xsig*xsig)
    term2 = (y-y0)**2 / (2*ysig*ysig)
    return height * np.exp(-(term1 + term2))

def moments(data):
    """ Returns (height, x, y, width_x, width_y) the gaussian
    parameters of a 2D distribution by calculating its moments."""
    total = data.sum()
    X, Y = np.indices(data.shape)
    x = (X*data).sum() / total
    y = (Y*data).sum() / total
    col = data[:, int(y)]
    width_x = np.sqrt(np.abs((np.arange(col.size)-y)**2*col).sum()/col.sum())
    row = data[int(x), :]
    width_y = np.sqrt(np.abs((np.arange(row.size)-x)**2*row).sum()/row.sum())
    height = data.max()
    return height, x, y, width_x, width_y

def fitgauss2d(x,y,data,sigma):
    """ Fit a 2d gaussian to data."""
    height, xc, yc, widthx, widthy = moments(data)
    xc,yc = x[int(xc)], y[int(yc)]
    xfwhm,yfwhm = widthy*np.diff(x).mean(), widthy*np.diff(y).mean()
    guess = dict(x0=xc,y0=yc,xfwhm=xfwhm,yfwhm=yfwhm,height=height)
    return astro.fit.fitfunc(np.meshgrid(x,y), data, sigma, gauss2d, guess)

if 1:
    from enthought.mayavi import mlab
    x = np.linspace(-15,15,100)
    y = np.linspace(-15,15,120)
    X,Y = np.meshgrid(x,y)
    gdata0  = gauss2d((X,Y),1,-1,8,4,20)
    gdata1 = gdata0 + np.random.randn(*gdata0.shape)
    mlab.mesh(X,Y,gdata1,representation='points',color=(1,0,0))
    mlab.axes()
    gsigma = np.ones(gdata1.shape)
    names,vals,m = fitgauss2d(x, y, gdata1, gsigma)
if 1:    
    # have a look at the results
    zip(names,vals,[m.errors[k] for k in names],(1,-1,8,4,20))
    pl.figure()
    pl.pcolor(np.array(m.matrix(correlation=True)),vmin=-1,vmax=1)
    pl.colorbar()

if 1:
    mlab.clf()
    mlab.mesh(X,Y,gdata1,representation='points',color=(1,0,0))
    mlab.mesh(X,Y,gauss2d((X,Y),*vals),opacity=0.7,color=(0.8,0.8,0.8))
    mlab.axes()
    
